# 代码功能说明：
# 这段代码演示了如何计算二维函数的数值梯度
# 实现了中心差分法计算梯度（支持批量处理）
# 可视化展示了函数f(x0,x1)=x0²+x1²的梯度场
# 梯度箭头指向函数值增加最快的方向（图中显示的是负梯度方向）
# 关键点：
# 使用中心差分法计算偏导数，比前向差分更精确
# 支持批量处理，可以同时计算多个点的梯度
# 使用matplotlib的quiver函数绘制梯度场
# 梯度箭头长度表示梯度大小，方向表示梯度方向

# 指定文件编码为UTF-8
# coding: utf-8

# 参考链接：http://d.hatena.ne.jp/white_wheels/20100327/p3
# 导入必要的库
import numpy as np
import matplotlib.pylab as plt
from mpl_toolkits.mplot3d import Axes3D  # 用于3D绘图

# 定义非批量的数值梯度计算函数
def _numerical_gradient_no_batch(f, x):
    h = 1e-4  # 微小的增量值
    grad = np.zeros_like(x)  # 创建与x形状相同的全零数组

    # 对x的每个元素计算偏导数
    for idx in range(x.size):
        tmp_val = x[idx]  # 保存原始值

        # 计算f(x+h)
        x[idx] = float(tmp_val) + h
        fxh1 = f(x)

        # 计算f(x-h)
        x[idx] = tmp_val - h
        fxh2 = f(x)

        # 计算中心差分梯度
        grad[idx] = (fxh1 - fxh2) / (2*h)

        x[idx] = tmp_val  # 恢复原始值

    return grad

# 定义数值梯度计算函数（支持批量处理）
def numerical_gradient(f, X):
    if X.ndim == 1:  # 如果是一维数组
        return _numerical_gradient_no_batch(f, X)
    else:  # 如果是二维数组（批量）
        grad = np.zeros_like(X)

        # 对每个样本计算梯度
        for idx, x in enumerate(X):
            grad[idx] = _numerical_gradient_no_batch(f, x)

        return grad

# 定义示例函数2：f(x0,x1) = x0^2 + x1^2
def function_2(x):
    if x.ndim == 1:  # 一维情况
        return np.sum(x**2)
    else:  # 二维情况（批量）
        return np.sum(x**2, axis=1)

# 定义切线函数（用于绘制梯度）
def tangent_line(f, x):
    d = numerical_gradient(f, x)  # 计算梯度
    print(d)  # 打印梯度值
    y = f(x) - d*x  # 计算切线截距
    return lambda t: d*t + y  # 返回切线函数

# 主程序入口
if __name__ == '__main__':
    # 创建x0和x1的网格点
    x0 = np.arange(-2, 2.5, 0.25)
    x1 = np.arange(-2, 2.5, 0.25)
    X, Y = np.meshgrid(x0, x1)  # 生成网格坐标

    # 将网格展平为一维数组
    X = X.flatten()
    Y = Y.flatten()

    # 计算各点的梯度
    grad = numerical_gradient(function_2, np.array([X, Y]))

    # 绘制梯度场
    plt.figure()
    # 使用箭头表示梯度方向（负梯度方向）
    plt.quiver(X, Y, -grad[0], -grad[1], angles="xy", color="#666666")
    plt.xlim([-2, 2])  # x轴范围
    plt.ylim([-2, 2])  # y轴范围
    plt.xlabel('x0')  # x轴标签
    plt.ylabel('x1')  # y轴标签
    plt.grid()  # 显示网格
    plt.legend()  # 显示图例
    plt.draw()  # 绘制图形
    plt.show()  # 显示图形
